#include <algorithm> // for swap :)

#include "LBM.h"

using std::swap;

void LatticeBoltzmannD2Q7::Visitor::walkCell()
{
  // propagation
  out().CR = CL().CR;
  out().UR = DL().UR;
  out().UL = DR().UL;
  out().CL = CR().CL;
  out().DL = UR().DL;
  out().DR = UL().DR;
  out().CC = CC().CC;

  if (ca.getCellType(outId()) == CELL_WALL)
  {
    swap(out().CR, out().CL);
    swap(out().UR, out().DL);
    swap(out().UL, out().DR);
    return;
  }

  LBfloat p = calcDensity(out());
  if (p < eps)
    return;

  LBfloat ux, uy;
  calcMomentum(out(), ux, uy);
  ux /= p;
  uy /= p;

  D2Q7Cell eq;
  ca.calcEquilibrium(p, ux, uy, eq);

  LBfloat coef = ca.getRelaxCoef();
  out().CR -= coef * (out().CR - eq.CR);
  out().UR -= coef * (out().UR - eq.UR);
  out().UL -= coef * (out().UL - eq.UL);
  out().CL -= coef * (out().CL - eq.CL);
  out().DL -= coef * (out().DL - eq.DL);
  out().DR -= coef * (out().DR - eq.DR);
  out().CC -= coef * (out().CC - eq.CC);
}

void LatticeBoltzmannD2Q7::calcEquilibrium(LBfloat p, LBfloat ux, LBfloat uy, D2Q7Cell & cell) const
{
  // equiilibrium distribution coefficients
  // f[i] = p*( r1 + r2*u^2 ) , i = 0
  // f[i] = p*( c1 + c2*(ei*u) + c3*(ei*u)^2 + c4*u^2 ) , i = 1..6
  static const LBfloat C1 = 1.0 / 12.0;
  static const LBfloat C2 = 1.0 / 3.0;
  static const LBfloat C3 = 2.0 / 3.0;
  static const LBfloat C4 = -1.0 / 6.0;

  static const LBfloat R1 = 1.0 / 2.0;
  static const LBfloat R2 = -1.0;

  LBfloat u2 = ux * ux + uy * uy;
  LBfloat ux1 = 0.5 * ux, uy1 = CellHeight * uy; 
  LBfloat sites[6] = {ux, ux1 + uy1, -ux1 + uy1, -ux, -ux1 - uy1, ux1 - uy1 };

  for (int i = 0; i != 6; ++i)
  {
    LBfloat eu = sites[i];
    LBfloat eu2 = eu * eu;
    sites[i] = p * (C1 + C2*eu + C3*eu2 + C4*u2);
  }

  cell.CC = p * (R1 + R2 * u2);
  cell.CR = sites[0];
  cell.UR = sites[1];
  cell.UL = sites[2];
  cell.CL = sites[3];
  cell.DL = sites[4];
  cell.DR = sites[5];
}


LatticeBoltzmannD2Q7::LatticeBoltzmannD2Q7(int width, int height)
  : Base(width, height), m_cellTypes(width, height)
{
  setViscosity(1.0);
}

void LatticeBoltzmannD2Q7::setConstBound(CellType cellType, const D2Q7Cell & cell)
{
  Base::setConstBound(cell);
  m_cellTypes.setConstBound(cellType);
}

void LatticeBoltzmannD2Q7::setTorusBound()
{
  Base::setTorusBound();
  m_cellTypes.setTorusBound();
}
